#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cusp/coo_matrix.h>
#include <cusp/gallery/poisson.h>
#include <cusp/blas.h>
#include <cusp/multiply.h>
#include <cusp/transpose.h>
#include <cusp/krylov/cg.h>
#include "FullWeighting.h"

typedef cusp::coo_matrix<int, float, cusp::host_memory> hMatrix; 
typedef cusp::coo_matrix<int, float, cusp::device_memory> dMatrix;
typedef cusp::array1d<float, cusp::host_memory> hVector; 
typedef cusp::array1d<float, cusp::device_memory> dVector;

struct MGLevel { 

	struct Level { 
		dMatrix A; 
		dVector b; 
		dVector x; 
		dMatrix restr; 
	}; 
	Level & operator[](int i) { 
		return levels[i]; 
	}; 
	MGLevel(int num) {
		n = num;
		levels = new Level[n+1];
	}
	
private: 
	Level * levels; 
	int n;
}; 

int
main(int argc, char * argv[])
{
	int nu1 = 2;
	int nu2 = 2;
	float omega = 0.8;
	int p = 8;
	int N = (1 << p) - 1;
	int minLevel = 3;
	int max = p - minLevel+1;
	MGLevel levels(max);
	
	hMatrix A_fine; 
	cusp::gallery::poisson5pt(A_fine, N, N);
	levels[0].A = A_fine;
	
	hVector x_fine(N*N); 
	levels[0].x = x_fine;
	
	hVector b_fine(N*N); cusp::blas::fill(b_fine, 1);
	levels[0].b = b_fine;
	
	for(int i=1; i < max; i++)
	{
		int Nfine = (1 << (p - i + 1) - 1);
		dVector r_fine(Nfine * Nfine);
		dVector y(Nfine * Nfine);
		
		//Relax
		for(int j=0; j<nu1; j++)
		{
			//r = b - A*x
			cusp::multiply(levels[i-1].A, levels[i-1].b, y);
			cusp::blas::axpby(levels[i-1].b, y, r_fine, 1, -1);
			//x = x + omega * r / 4;
			cusp::blas::axpy(r_fine, levels[i-1].x, omega / 4);
		}
		
		//r = b - A*x
		cusp::multiply(levels[i-1].A, levels[i-1].x, y);
		cusp::blas::axpby(levels[i-1].b, y, r_fine, 1, -1);
		
		//Prepare data for next level
		int Ni = (1 << (p - i) - 1 );
		dVector r_coarse(Ni * Ni);
		
		//Restrict residual
		hMatrix restr;
		full_weighting_2D(restr, Ni);
		levels[i-1].restr = restr;
		cusp::multiply(levels[i-1].restr, r_fine, r_coarse);
		
		hMatrix A;
		cusp::gallery::poisson5pt(A, Ni, Ni);
		levels[i].A = A;
		
		cusp::blas::copy(r_coarse,levels[i].b);
		
		hVector x(Ni * Ni); cusp::blas::fill(x, 0);
		cusp::blas::copy(x, levels[i].x);
	}
	
	//Coarsest level
	cusp::krylov::cg(levels[max].A, levels[max].x, levels[max].b);

	for(int i=max-1; i >= 0; i--)
	{   
		//Pad x with zeros
		int Ni = (1<< i) - 1;
		dMatrix interp;
		//cusp::transpose(levels[i+1].restr, interp);
		
/*		//Interpolate
		dVector r(Ni * Ni);
		cusp::multiply(interp, levels[i+1].x, r);
		
		//Correct
		cusp::blas::axpy(r, levels[i].x, 1);
		
		dVector y(Ni * Ni);
		//Relax
		for(int j=0; j<nu2; j++)
		{
			//r = b - A*x
			cusp::multiply(levels[i].A, levels[i].x, y);
			cusp::blas::axpby(levels[i].b, y, r, 1, -1);
			//x = x + omega * r / 4;
			cusp::blas::axpy(r, levels[i].x, omega / 4);
		}*/
	}

	return 0;
}
